import sqlite3
import os
import datetime

class DatabaseManager:
    def __init__(self, db_path):
        self.db_path = db_path
        self.conn = None
        self.create_tables()

    def reset(self):
        with self as cursor:
            cursor.execute('DROP TABLE IF EXISTS scores')
            self.create_tables()

    def create_tables(self):
        with self as cursor:
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS scores (
                    grade TEXT NOT NULL,
                    exam_id TEXT NOT NULL,
                    class_id TEXT NOT NULL,
                    student_id TEXT NOT NULL,
                    name TEXT NOT NULL,
                    biology_score REAL,
                    school_rank INTEGER,
                    class_rank INTEGER
                )
            ''')

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        return self.conn.cursor()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.conn:
            if exc_type is None:
                self.conn.commit()
            else:
                self.conn.rollback()
            self.conn.close()
        self.conn = None  # 彻底关闭连接

    def execute(self, query, params=None):
        with self as cursor:
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)
            return cursor

    def backup(self, backup_dir='backup'):
        if not os.path.exists(backup_dir):
            os.makedirs(backup_dir)
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        backup_path = os.path.join(backup_dir, f'backup_{timestamp}.db')
        with self.conn:
            self.conn.backup(sqlite3.connect(backup_path))
        return backup_path

    def restore(self, backup_path):
        if not os.path.exists(backup_path):
            raise FileNotFoundError(f"备份文件 {backup_path} 不存在")
        
        # 关闭原有连接
        if self.conn:
            self.conn.close()
        
        # 复制备份文件覆盖当前数据库
        import shutil
        shutil.copy(backup_path, self.db_path)
        
        # 重新初始化连接
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.create_tables()

    def get_all_scores(self):
        with self as cursor:
            cursor.execute('''SELECT grade, exam_id, class_id, student_id, name, biology_score, school_rank, class_rank 
                FROM scores''')
            columns = [col[0] for col in cursor.description]
            return [dict(zip(columns, row)) for row in cursor.fetchall()]

    def import_scores(self, records):
        with self as cursor:
            cursor.executemany('''
                INSERT INTO scores (
                    grade, exam_id, class_id, student_id, name,
                    biology_score, school_rank, class_rank
                ) VALUES (
                    :grade, :exam_id, :class_id, :student_id, :name,
                    :biology_score, :school_rank, :class_rank
                )
            ''', records)

    def get_database_status(self):
        with self as cursor:
            cursor.execute('SELECT COUNT(DISTINCT exam_id) FROM scores')
            exam_count = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(DISTINCT class_id) FROM scores')
            class_count = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(DISTINCT student_id) FROM scores')
            student_count = cursor.fetchone()[0]
            
            return {
                'exam_count': exam_count,
                'class_count': class_count,
                'student_count': student_count
            }
            
    def get_all_classes(self):
        with self as cursor:
            cursor.execute('SELECT DISTINCT class_id FROM scores ORDER BY class_id')
            return [row[0] for row in cursor.fetchall()]

    def get_all_grade(self):
        with self as cursor:
            cursor.execute('SELECT DISTINCT grade FROM scores ORDER BY grade')
            return [row[0] for row in cursor.fetchall()]

    def get_all_exams(self):
        with self as cursor:
            cursor.execute('SELECT DISTINCT exam_id FROM scores ORDER BY exam_id')
            return [row[0] for row in cursor.fetchall()]

    def conditional_query(self, grade=None, exam_id=None, class_id=None, name=None):
        query = 'SELECT * FROM scores WHERE 1=1'
        params = []
        
        if grade:
            query += ' AND grade = ?'
            params.append(grade)
        if exam_id:
            query += ' AND exam_id = ?'
            params.append(exam_id)
        if class_id:
            query += ' AND class_id = ?'
            params.append(class_id)
        if name:
            query += ' AND name LIKE ?'
            params.append(f'%{name}%')
        
        with self as cursor:
            cursor.execute(query, params)
            columns = [col[0] for col in cursor.description]
            return [dict(zip(columns, row)) for row in cursor.fetchall()]

    def get_students_names(self, exam_id, class_id):
        with self.conn:
            cursor = self.conn.cursor()
            cursor.execute('''
                SELECT name FROM scores 
                WHERE exam_id = ? AND class_id = ?
            ''', (exam_id, class_id))
            return [row[0] for row in cursor.fetchall()]